# -*- coding: UTF-8 -*-
import wave
import os
import numpy as np
import scipy.io.wavfile as wav
import matplotlib.pyplot as plt


def sgn(data):
    if data >= 0:
        return 1
    else:
        return -1


cai_num = 256


# 计算每一帧的能量 256个采样点为一帧
def calEnergy(wave_data):
    energy = []
    sum = 0
    for i in range(len(wave_data)):
        sum = sum + (int(wave_data[i]) * int(wave_data[i]))
        if (i + 1) % cai_num == 0:
            energy.append(sum)
            sum = 0
        elif i == len(wave_data) - 1:
            energy.append(sum)
    return energy


# 计算过零率
def calZeroCrossingRate(wave_data):
    zeroCrossingRate = []
    sum = 0
    for i in range(len(wave_data)):
        if i % cai_num == 0:
            continue
        sum = sum + np.abs(sgn(wave_data[i]) - sgn(wave_data[i - 1]))
        if (i + 1) % cai_num == 0:
            zeroCrossingRate.append(float(sum) / (cai_num-1))
            sum = 0
        elif i == len(wave_data) - 1:
            zeroCrossingRate.append(float(sum) / (cai_num-1))
    return zeroCrossingRate


# 利用短时能量，短时过零率，使用双门限法进行端点检测
def endPointDetect(wave_data, energy, zeroCrossingRate):
    sum = 0
    energyAverage = 0
    for en in energy:
        sum = sum + en
    energyAverage = sum / len(energy)

    sum = 0
    for en in energy[:5]:
        sum = sum + en
    ML = sum / 5
    MH = energyAverage / 4  # 较高的能量阈值
    ML = (ML + MH) / 4  # 较低的能量阈值
    sum = 0
    for zcr in zeroCrossingRate[:5]:
        sum = float(sum) + zcr
    Zs = sum / 5  # 过零率阈值

    A = []
    B = []
    C = []

    # 首先利用较大能量阈值 MH 进行初步检测
    flag = 0
    for i in range(len(energy)):
        if len(A) == 0 and flag == 0 and energy[i] > MH:
            A.append(i)
            flag = 1
        elif flag == 0 and energy[i] > MH and i - 21 > A[len(A) - 1]:
            A.append(i)
            flag = 1
        elif flag == 0 and energy[i] > MH and i - 21 <= A[len(A) - 1]:
            A = A[:len(A) - 1]
            flag = 1

        if flag == 1 and energy[i] < MH:
            A.append(i)
            flag = 0
    # print("较高能量阈值，计算后的浊音A:" + str(A))

    # 利用较小能量阈值 ML 进行第二步能量检测
    for j in range(len(A)):
        i = A[j]
        if j % 2 == 1:
            while i < len(energy) and energy[i] > ML:
                i = i + 1
            B.append(i)
        else:
            while i > 0 and energy[i] > ML:
                i = i - 1
            B.append(i)
    # print("较低能量阈值，增加一段语言B:" + str(B))

    # 利用过零率进行最后一步检测
    for j in range(len(B)):
        i = B[j]
        if j % 2 == 1:
            while i < len(zeroCrossingRate) and zeroCrossingRate[i] >= 3 * Zs:
                i = i + 1
            C.append(i)
        else:
            while i > 0 and zeroCrossingRate[i] >= 3 * Zs:
                i = i - 1
            C.append(i)
    # print("过零率阈值，最终语音分段C:" + str(C))
    return C


# 双门限法 端点检测处理
# 文件处理
def endPoint(inputPath, outPath):
    fs1, wave_data = wav.read(inputPath)
    # print(fs1)
    # print(wave_data)
    # plt.rcParams["font.sans-serif"] = ["SimHei"]
    # plt.rcParams["axes.unicode_minus"] = False
    # plt.plot(wave_data)
    # plt.title("音频信号")
    # plt.show()
    # 计算短时能量
    energy = calEnergy(wave_data)
    # print("短时能量：", energy)
    # plt.plot(energy)
    # plt.title("短时能量")
    # plt.show()
    # 计算短时过零率
    rate = calZeroCrossingRate(wave_data)
    # print("短时过零率：", rate)
    # plt.plot(rate)
    # plt.title("短时过零率")
    # plt.show()
    # 进行端点检验
    detect = endPointDetect(wave_data, energy, rate)
    # print("端点分段：", detect)
    detect = [i * cai_num for i in detect]
    # print("音频分段：", detect)
    if len(detect) == 1:
        detect.append(len(wave_data))
    endWaveFile = wave_data[detect[0]:detect[1]]
    # plt.plot(endWaveFile)
    # plt.title("语音部分")
    # plt.show()
    # output = "E:\\DeepLearning\\dl_studio\\bishe\\Audio\\MyAudioData\\test1_output.wav"
    wav.write(outPath, fs1, endWaveFile)
    print("保存成功：", outPath)


# 返回一维list
def endPointReTurnNp(wave_data):
    # fs1, wave_data = wav.read(inputPath)
    # print(fs1)
    # print(wave_data)
    # plt.rcParams["font.sans-serif"] = ["SimHei"]
    # plt.rcParams["axes.unicode_minus"] = False
    # plt.plot(wave_data)
    # plt.title("音频信号")
    # plt.show()
    # 计算短时能量
    energy = calEnergy(wave_data)
    # print("短时能量：", energy)
    # plt.plot(energy)
    # plt.title("短时能量")
    # plt.show()
    # 计算短时过零率
    rate = calZeroCrossingRate(wave_data)
    # print("短时过零率：", rate)
    # plt.plot(rate)
    # plt.title("短时过零率")
    # plt.show()
    # 进行端点检验
    detect = endPointDetect(wave_data, energy, rate)
    # print("端点分段：", detect)
    detect = [i * cai_num for i in detect]
    # print("音频分段：", detect)
    if len(detect) == 1:
        detect.append(len(wave_data))
    # 循环把有音频的部分全部连接起来
    endWaveFile = []
    for i in range(len(detect) - 1):
        # i: 1 2 3
        # i+1: 2 3 4
        # [192768, 269312, 277248, 279552, 285696, 326400, 604928, 605696]
        endWaveFile.extend(wave_data[detect[i]:detect[i + 1]])
    # endWaveFile = wave_data[detect[0]:detect[1]]
    # plt.plot(endWaveFile)
    # plt.title("语音部分")
    # plt.show()
    # output = "E:\\DeepLearning\\dl_studio\\bishe\\Audio\\MyAudioData\\test1_output.wav"
    # print(endWaveFile)
    return endWaveFile


# for i in range(10):
#     f = wave.open("./语料/" + str(i + 1) + ".wav", "rb")
#     # getparams() 一次性返回所有的WAV文件的格式信息
#     params = f.getparams()
#     # nframes 采样点数目
#     nchannels, sampwidth, framerate, nframes = params[:4]
#     # readframes() 按照采样点读取数据
#     str_data = f.readframes(nframes)  # str_data 是二进制字符串
#
#     # 以上可以直接写成 str_data = f.readframes(f.getnframes())
#
#     # 转成二字节数组形式（每个采样点占两个字节）
#     wave_data = np.fromstring(str_data, dtype=np.short)
#     print("采样点数目：" + str(len(wave_data)))  # 输出应为采样点数目
#     f.close()
#     energy = calEnergy(wave_data)
#     with open("./energy/" + str(i + 1) + "_en.txt", "w") as f:
#         for en in energy:
#             f.write(str(en) + "\n")
#     zeroCrossingRate = calZeroCrossingRate(wave_data)
#     with open("./zeroCrossingRate/" + str(i + 1) + "_zero.txt", "w") as f:
#         for zcr in zeroCrossingRate:
#             f.write(str(zcr) + "\n")
#     N = endPointDetect(wave_data, energy, zeroCrossingRate)
#     # 输出为 pcm 格式
#     with open("./端点检测后的语料/" + str(i + 1) + ".pcm", "wb") as f:
#         i = 0
#         while i < len(N):
#             for num in wave_data[N[i] * 256: N[i + 1] * 256]:
#                 f.write(num)
#             i = i + 2
if __name__ == '__main__':
    test1 = "E:\\DeepLearning\\dl_studio\\bishe\\Audio\\MyAudioData\\test1.wav"
    fs1, wave_data = wav.read(test1)
    print(fs1)
    print(wave_data)
    #
    print(len(wave_data) / fs1)
    plt.rcParams["font.sans-serif"] = ["SimHei"]
    plt.rcParams["axes.unicode_minus"] = False
    plt.plot(wave_data)
    plt.title("音频信号")
    plt.show()
    # 计算短时能量
    energy = calEnergy(wave_data)
    print("短时能量：", energy)
    plt.plot(energy)
    plt.title("短时能量")
    plt.show()
    # 计算短时过零率
    rate = calZeroCrossingRate(wave_data)
    print("短时过零率：", rate)
    plt.plot(rate)
    plt.title("短时过零率")
    plt.show()
    # 进行端点检验
    detect = endPointDetect(wave_data, energy, rate)
    print("端点分段：", detect)
    detect = [i * cai_num for i in detect]
    print("音频分段：", detect)
    endWaveFile = wave_data[detect[0]:detect[1]]
    plt.plot(endWaveFile)
    plt.title("语音部分")
    plt.show()
    output = "E:\\DeepLearning\\dl_studio\\bishe\\Audio\\MyAudioData\\test1_output.wav"
    wav.write(output, fs1, endWaveFile)
    print("=============保存成功==============")
